package com.xwq.xwqaiagent.rag;

import com.alibaba.cloud.ai.reader.mysql.MySQLDocumentReader;
import com.alibaba.cloud.ai.reader.mysql.MySQLResource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * ai恋爱大师文档读取器
 */
@Component
@Slf4j
class LoveAppDocumentLoader {

    private final ResourcePatternResolver resourcePatternResolver;


    LoveAppDocumentLoader(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }

    /**
     * 读取并加载Md文件
     * @return
     */
    public List<Document> loadMarkdowns() {
        List<Document> allDocuments = new ArrayList<>();
        try {
            // 这里可以修改为你要加载的多个 Markdown 文件的路径模式
            Resource[] resources = resourcePatternResolver.getResources("classpath:document/*.md");
            for (Resource resource : resources) {
                String fileName = resource.getFilename();
                // 提取文档倒数第 3 和第 2 个字作为标签
                //String status = fileName.substring(fileName.length() - 6, fileName.length() - 4);
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(false)
                        .withIncludeBlockquote(false)
                        .withAdditionalMetadata("filename", fileName)
                        //.withAdditionalMetadata("status", status)
                        .build();

                MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
                allDocuments.addAll(reader.get());
            }
        } catch (IOException e) {
            log.error("Markdown 文档加载失败", e);
        }
        return allDocuments;
    }

    /**
     * 读取并加载 MySQL 数据库中的数据
     * @return 包含数据库记录的文档列表
     */
    public List<Document> loadMysql() {
        List<Document> allDocuments = new ArrayList<>();
        try {
            // 创建 MySQL 资源（从配置文件中读取）
            String host = "localhost"; // 从配置文件中读取
            int port = 3306;
            String database = "xwq_agent";
            String username = "root";
            String password = "1234";
            String query = "SELECT * FROM chat_message WHERE is_delete = 0 LIMIT 5";
            List<String> contentColumns = Arrays.asList("content", "message_type", "conversation_id");
            List<String> metadataColumns = Arrays.asList("id", "create_time", "update_time");

            MySQLResource mysqlResource = new MySQLResource(
                    host, port, database, username, password, query, contentColumns, metadataColumns
            );

            MySQLDocumentReader reader = new MySQLDocumentReader(mysqlResource);
            allDocuments.addAll(reader.get());
        } catch (Exception e) {
            log.error("MySQL 数据加载失败", e);
        }
        return allDocuments;
    }
}
